from dotenv import load_dotenv, find_dotenv
_ = load_dotenv(find_dotenv())

# 确定选择使用的模型服务平台
# supplier = "OpenAI"
supplier = "ZhipuAI"

# 根据服务平台初始化模型名称及API
if supplier == "OpenAI":
    from openai import OpenAI
    client = OpenAI()
    model = "dall-e-3"
    functionName = client.images.generate
else:
    from zhipuai import ZhipuAI
    client = ZhipuAI()
    model = "cogview-3-plus"
    functionName = client.images.generations

def get_completion(prompt, model="cogview-3-plus", size="1024x1024"):
    response = functionName(
      model=model, 
      prompt=prompt,
      size=size
    )
    return response.data[0].url
  
while True:  
    user_input = input("请输入待生成的图片描述（输入'exit'退出程序）: ")  
    if user_input is None or user_input == "" :
        continue
    elif user_input.lower() == 'exit':  
        print("程序退出。\n")  
        break  
    else :
        res = get_completion(prompt=user_input, model=model)  
        print(f"图片URL：{res}")
